#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
NLTK数据下载对话框模块
提供界面用于下载NLTK数据
"""

import os
import sys
import logging
from pathlib import Path
from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, 
    QProgressBar, QListWidget, QMessageBox
)
from PyQt5.QtCore import pyqtSignal, QThread, Qt

try:
    import nltk
    nltk_available = True
except ImportError:
    nltk_available = False

# 导入项目模块
sys.path.append(str(Path(__file__).resolve().parent.parent))
from config.paths import NLTK_DATA_DIR, create_directories
from core.model_manager import ModelManager

# 配置日志
logger = logging.getLogger("NLTK下载")

class NLTKDownloadThread(QThread):
    """NLTK数据下载线程"""
    
    progress = pyqtSignal(str)
    finished = pyqtSignal(bool)
    
    def __init__(self):
        """初始化下载线程"""
        super().__init__()
        self.model_manager = ModelManager(ui_callback=self.log_callback)
    
    def log_callback(self, message, is_error=False):
        """发送日志消息到UI"""
        self.progress.emit(message)
    
    def run(self):
        """线程主函数"""
        try:
            # 下载NLTK数据
            self.progress.emit("开始下载NLTK数据...")
            
            # 配置NLTK只使用项目内目录
            os.environ['NLTK_DATA'] = str(NLTK_DATA_DIR)
            nltk.data.path = []  # 清空当前路径
            nltk.data.path.append(str(NLTK_DATA_DIR))  # 只使用项目内目录
            
            # 下载所需数据包
            self.progress.emit("下载punkt数据包...")
            nltk.download('punkt', download_dir=str(NLTK_DATA_DIR))
            
            self.progress.emit("下载averaged_perceptron_tagger数据包...")
            nltk.download('averaged_perceptron_tagger', download_dir=str(NLTK_DATA_DIR))
            
            self.progress.emit("NLTK数据下载完成")
            self.finished.emit(True)
        except Exception as e:
            self.progress.emit(f"下载过程中出错: {str(e)}")
            self.finished.emit(False)


class NLTKDownloadDialog(QDialog):
    """NLTK数据下载对话框"""
    
    def __init__(self, parent=None):
        """初始化对话框"""
        super().__init__(parent)
        
        # 确保必要的目录结构存在
        create_directories()
        
        self.setWindowTitle("NLTK数据下载")
        self.setMinimumSize(500, 300)
        
        # 初始化界面
        self.init_ui()
    
    def init_ui(self):
        """初始化用户界面"""
        # 主布局
        layout = QVBoxLayout(self)
        
        # 说明标签
        info_label = QLabel(
            "需要下载NLTK数据才能使用文本到语音功能。\n"
            "点击下载按钮开始自动下载所需数据。"
        )
        info_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(info_label)
        
        # 下载按钮
        self.download_button = QPushButton("下载NLTK数据")
        self.download_button.clicked.connect(self.start_download)
        layout.addWidget(self.download_button)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(True)
        self.progress_bar.setValue(0)
        self.progress_bar.setVisible(False)
        layout.addWidget(self.progress_bar)
        
        # 日志区域
        self.log_list = QListWidget()
        layout.addWidget(self.log_list)
        
        # 关闭按钮
        self.close_button = QPushButton("关闭")
        self.close_button.clicked.connect(self.accept)
        layout.addWidget(self.close_button)
    
    def log_message(self, message):
        """添加日志消息"""
        self.log_list.addItem(message)
        self.log_list.scrollToBottom()
    
    def start_download(self):
        """开始下载NLTK数据"""
        if not nltk_available:
            QMessageBox.warning(
                self, 
                "NLTK未安装", 
                "NLTK未正确安装，请先安装NLTK库"
            )
            return
        
        # 设置UI状态
        self.download_button.setEnabled(False)
        self.close_button.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setRange(0, 0)  # 设置为忙碌状态
        
        # 创建并启动下载线程
        self.download_thread = NLTKDownloadThread()
        self.download_thread.progress.connect(self.log_message)
        self.download_thread.finished.connect(self.on_download_finished)
        self.download_thread.start()
    
    def on_download_finished(self, success):
        """
        下载完成的回调
        
        Args:
            success: 是否成功
        """
        # 恢复UI状态
        self.download_button.setEnabled(True)
        self.close_button.setEnabled(True)
        self.progress_bar.setRange(0, 100)
        
        if success:
            self.progress_bar.setValue(100)
            QMessageBox.information(
                self, 
                "下载完成", 
                "NLTK数据下载和配置完成！"
            )
        else:
            self.progress_bar.setValue(0)
            QMessageBox.warning(
                self, 
                "下载失败", 
                "NLTK数据下载或配置失败，请检查网络连接后重试。"
            )


# 测试代码
if __name__ == "__main__":
    from PyQt5.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    dialog = NLTKDownloadDialog()
    dialog.show()
    sys.exit(app.exec_()) 